{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/src/MaxText/examples/multimodal_gemma3_demo.ipynb)\n",
    "\n",
    "# Gemma3 Multimodal Inference/Training Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This notebook demonstrates MaxText's multimodal features, using Gemma3-4B as an example:\n",
    "- Convert an orbax checkpoint from HuggingFace.\n",
    "- Apply decoding on a single image input.\n",
    "- Apply SFT to the converted checkpoint on ChartQA dataset.\n",
    "\n",
    "Given the relative small size of Gemma3-4B, you can run this colab on a v4-8, v5p-8 or v6e-4 TPU VM. You can also use [XPK](https://github.com/AI-Hypercomputer/maxtext/blob/64d6d9b425e78dde94c37a82bb13ba5606e74b1b/docs/guides/run_maxtext_via_xpk.md) to run training workloads on a TPU cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Your Hugging Face Token\n",
    "\n",
    "To access model checkpoint from the Hugging Face Hub, you need to authenticate with a personal access token.\n",
    "\n",
    "**Follow these steps to get your token:**\n",
    "\n",
    "1.  **Navigate to the Access Tokens page** in your Hugging Face account settings. You can go there directly by visiting this URL:\n",
    "    *   [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)\n",
    "\n",
    "2.  **Create a new token** by clicking the **\"+ Create new token\"** button.\n",
    "\n",
    "3.  **Give your token a name** and assign it a **`read` role**. The `read` role is sufficient for downloading models.\n",
    "\n",
    "4.  **Copy the generated token**. You will need to paste it in `HF_TOKEN`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5KPyOE8e9WbO"
   },
   "outputs": [],
   "source": [
    "# Install maxtext and dependencies\n",
    "# 1. Install uv, a fast Python package installer\n",
    "!pip install uv\n",
    "\n",
    "# 2. Install MaxText and its dependencies\n",
    "!uv pip install maxtext --resolution=lowest\n",
    "!install_maxtext_github_deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import MaxText\n",
    "\n",
    "# Get the root directory of the MaxText\n",
    "MAXTEXT_REPO_ROOT = os.path.dirname(MaxText.__file__)\n",
    "\n",
    "# Define model name\n",
    "MODEL_NAME = \"gemma3-4b\"\n",
    "\n",
    "# Use either a GCS path or a local path for the model checkpoint\n",
    "MODEL_CHECKPOINT_PATH = f\"gs://your-gcs-bucket/{MODEL_NAME}\"\n",
    "\n",
    "# Replace with your actual Hugging Face token\n",
    "HF_TOKEN = \"your_huggingface_token_here\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert Checkpoint from HuggingFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m MaxText.utils.ckpt_conversion.to_maxtext \\\n",
    "    $MAXTEXT_REPO_ROOT/configs/base.yml \\\n",
    "    model_name=$MODEL_NAME \\\n",
    "    hf_access_token=$HF_TOKEN \\\n",
    "    base_output_directory=$MODEL_CHECKPOINT_PATH \\\n",
    "    use_multimodal=true \\\n",
    "    scan_layers=false"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decode on One Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m MaxText.decode \\\n",
    "    $MAXTEXT_REPO_ROOT/configs/base.yml \\\n",
    "    model_name=$MODEL_NAME \\\n",
    "    tokenizer_path=assets/tokenizer.gemma3 \\\n",
    "    load_parameters_path=$MODEL_CHECKPOINT_PATH/0/items \\\n",
    "    per_device_batch_size=1 \\\n",
    "    run_name=ht_test max_prefill_predict_length=272 \\\n",
    "    max_target_length=300 \\\n",
    "    steps=1 \\\n",
    "    async_checkpointing=false \\\n",
    "    scan_layers=false \\\n",
    "    use_multimodal=true \\\n",
    "    prompt='Describe image <start_of_image>' \\\n",
    "    image_path=$MAXTEXT_REPO_ROOT/test_assets/test_image.jpg \\\n",
    "    attention='dot_product'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supervised Finetuning (SFT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running the cell below will trigger a 10-step SFT on your TPU VM (v4-8, v5p-8, or v6e-4). However, we recommend using [XPK](https://github.com/AI-Hypercomputer/maxtext/blob/64d6d9b425e78dde94c37a82bb13ba5606e74b1b/docs/guides/run_maxtext_via_xpk.md) to schedule a training workload on a TPU cluster for better performance. After the SFT, the result checkpoint will be saved to `BASE_OUTPUT_DIRECTORY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define SFT output directory\n",
    "BASE_OUTPUT_DIRECTORY=f\"gs://your-gcs-bucket/{MODEL_NAME}-sft\"\n",
    "PRE_TRAINED_MODEL_TOKENIZER=\"google/gemma-3-4b-it\"\n",
    "WORKLOAD_NAME=f\"{MODEL_NAME}-chartqa-sft\"\n",
    "STEPS=10\n",
    "PER_DEVICE_BATCH_SIZE=1\n",
    "\n",
    "!python -m MaxText.sft_trainer \\\n",
    "    $MAXTEXT_REPO_ROOT/configs/sft-vision-chartqa.yml \\\n",
    "    run_name=$WORKLOAD_NAME \\\n",
    "    model_name=$MODEL_NAME \\\n",
    "    tokenizer_path=$PRE_TRAINED_MODEL_TOKENIZER \\\n",
    "    hf_access_token=$HF_TOKEN \\\n",
    "    load_parameters_path=$MODEL_CHECKPOINT_PATH/0/items \\\n",
    "    base_output_directory=$BASE_OUTPUT_DIRECTORY \\\n",
    "    per_device_batch_size=$PER_DEVICE_BATCH_SIZE \\\n",
    "    steps=$STEPS \\\n",
    "    max_prefill_predict_length=1024 \\\n",
    "    max_target_length=2048 \\\n",
    "    checkpoint_period=1000 \\\n",
    "    scan_layers=False \\\n",
    "    async_checkpointing=True \\\n",
    "    enable_checkpointing=True \\\n",
    "    attention=dot_product \\\n",
    "    max_num_images_per_example=1 \\\n",
    "    dataset_type=hf profiler=xplane"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "gpuType": "V5E1",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "python3.12",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
